import ptan
import numpy as np

import torch
import torch.distributions.constraints
import torch.nn as nn
import torch.nn.functional as F
import torch.distributions as dist
from typing import List, Optional
import einops as eop
from conv_kan.efficient_kan import KAN as EFF_KAN
from conv_kan.KANLinear import KAN as LIN_KAN
from torch.distributions import Normal, Independent
import lib.common as common

HID_SIZE = 128


'''
在强化学习的Actor-Critic方法中，特别是在处理连续动作空间时，网络通常预测动作的均值（`mu`）和方差（`var`）。在你给出的代码中，模型使用不同的激活函数来处理这两个输出，即`tanh`用于均值，而`Softplus`用于方差。让我们分别来看看这两个激活函数的选择原因：

### 1. 均值 (`mu`) 的 `tanh` 激活函数

- **动作范围限制**：`tanh` 函数的输出范围在 [-1, 1] 之间。在许多连续动作空间的问题中，动作的有效范围被限制在这个区间内。例如，如果你在一个物理模拟环境中控制机械臂，动作可能是旋转角度，这些角度通常有最大和最小值。

- **输出规范化**：使用 `tanh` 确保网络输出的均值永远不会超出这个范围，这样可以保证生成的动作始终在有效的动作空间内。

### 2. 方差 (`var`) 的 `Softplus` 激活函数

- **保证正值**：方差是衡量数据离散程度的度量，它必须是正值。`Softplus` 函数是一个平滑的函数，其输出始终为正值。`Softplus` 函数的公式是 `log(1 + exp(x))`，对于所有的输入 `x`，输出都是正的。

- **灵活性**：与 `ReLU` 函数相比，`Softplus` 在 `x` 接近0时更加平滑，这允

许更精细的控制方差的值。`Softplus` 提供了一种有效的方式来表示方差，即使在其原始预测值接近零或负值时也能保持正数。这在学习过程中非常有用，特别是当网络还未充分训练，对方差的估计可能接近零或变化较大时。

### 结合使用 `tanh` 和 `Softplus`

在这个A2C模型中，`tanh` 用于输出规范化到特定范围内的动作均值，而 `Softplus` 用于确保方差始终为正值且平滑。这种设计使模型能够有效地操作连续动作空间，同时保持必要的灵活性和稳定性。使用这两个激活函数的组合，模型可以生成符合实际物理限制的动作（如限制在一定范围内的力或速度），并且以合理的方式探索这些动作的不确定性。
'''
class ModelA2C(nn.Module):
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 观测的数据维度，不是图像数据，所以后续用的是全连接层
        act_size: 动作空间的维度，在这个游戏里面，指的是同时执行动作的数量
        '''

        super(ModelA2C, self).__init__()

        # 只有这个是提取特征，其余的都是输出结果
        self.base = nn.Sequential(
            nn.Linear(obs_size, HID_SIZE),
            nn.ReLU(),
        )
        # 输出均值，表示最终要执行的动作内容
        # 这里由于预测的输出动作包含负值，所以采用tanh函数，将输出值限制在-1到1之间
        # 而不是用sigmoid
        self.mu = nn.Sequential(
            nn.Linear(HID_SIZE, act_size),
            nn.Tanh(),
        )
        # https: // zhuanlan.zhihu.com / p / 461707201
        # var作用 方差平方，怀疑是用来稳定输出的概率范围大小
        # 用来促进网络进行探索以及指导网路朝哪个方向进行训练，使得整体趋近中值，但是在这里并没有直接使用方差，而是使用了信息熵的方式
        self.var = nn.Sequential(
            nn.Linear(HID_SIZE, act_size),
            nn.Softplus(), # Relu的替代函数，用于解决梯度消失问题 具体使用场景查看笔记内容
        )
        # 状态值（Q值），用来评价当前Q值，来评估当前执行的动作是否有优势
        self.value = nn.Linear(HID_SIZE, 1)

    def forward(self, x):
        '''
        return 均值，方差平方，Q值
        '''
        base_out = self.base(x)
        return self.mu(base_out), self.var(base_out), self.value(base_out)


class SACActor(nn.Module):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, obs_size, act_size, hidden_dim, action_range = 1., init_w = 3e-3, log_std_min = -20, log_std_max = 2):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(SACActor, self).__init__()
        self.log_std_min = log_std_min
        self.log_std_max = log_std_max
        self.action_range = action_range
        self.action_range_tensor = torch.tensor(action_range, dtype=torch.float32)
        self.act_size = act_size

        # sac代码中对网络权重进行了初始化，这里是对网络权重进行初始化，是否是必须的？
        # 经过试验不是必须要的
        self.net = nn.Sequential(
            nn.Linear(obs_size, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU()
        )

        # 和TD3的网络结构不同点，输出的是均值和对数标准差
        # 为什么采用对数标准差：
        # 在强化学习的上下文中，特别是在使用类似Soft Actor-Critic (SAC) 这样的算法时，`log_std`是指对数标准差，即动作分布的标准差的对数。这种表示方法在深度学习和统计模型中很常见，特别是在参数化某些概率分布（如正态分布）时。
        #
        # ### **为什么使用对数标准差（log_std）?**
        #
        # 使用对数标准差而不是直接使用标准差（std）有几个好处：
        #
        # 1. **数值稳定性**：通过将标准差的对数作为优化参数，可以确保在梯度下降过程中标准差保持正数。标准差必须是正数，而直接优化标准差可能在更新过程中导致其变为零或负数，这在数学上没有意义并可能导致运行时错误。
        #
        # 2. **无限制的值范围**：对数标准差可以取任意实数值（正无穷到负无穷），这使得优化过程更灵活和容易。对数转换后，参数的范围变得无约束，这对使用诸如梯度下降之类的优化算法更为方便。，而标准差只能是正数，所以采用对数标准差
        #
        # 3. **改善学习动态**：对数变换能够扩展参数的有效学习范围，使得小的变化也能在梯度下降中有效反映，增强了模型在参数空间的探索能力。
        #
        # ### **在策略网络中的应用**
        #
        # 在SAC等算法中，策略网络通常输出一个动作的均值（mean）和对数标准差（log_std）。这些输出被用来参数化动作的概率分布，通常是一个多元正态分布。策略网络的输出如下：
        #
        # - **均值（mean）**：描述了给定状态下预期动作的中心位置。
        # - **对数标准差（log_std）**：描述了动作取值的离散程度，高标准差意味着高不确定性和探索性，低标准差则表示动作更加确定。
        #
        # 在从策略中采样动作时，通常的步骤是首先从网络获取均值和对数标准差，计算标准差（通过指数化log_std），然后从以均值为中心、计算得到的标准差为离散度的正态分布中采样。使用`tanh`等激活函数可以进一步将动作限制在特定范围内（如[-1, 1]）。
        #
        # 总的来说，`log_std`在算法中的使用增强了模型的灵活性和稳定性，同时支持了有效的探索机制，这对于解决复杂的连续动作空间问题至关重要。
        self.mean_linear = nn.Linear(hidden_dim, act_size)
        self.log_std_linear = nn.Linear(hidden_dim, act_size)

    def forward(self, x):
        x = self.net(x)
        mean = self.mean_linear(x)
        # 限制标准差的范围，作用如下：
        # 避免极端行为、数值稳定性、保持有效的探索、便于调参
        log_std = self.log_std_linear(x)
        log_std = torch.clamp(log_std, self.log_std_min, self.log_std_max)

        return mean, log_std


    def evaluate(self, state, device, epsilon=1e-6):
        '''

        :param state:
        :param device:
        :param epsilon:
        :return: 执行的动作，动作的对数概率，动作噪声（动作的采样），动作均值，动作的对数标准差
        '''

        state_v = state
        # 得到预测动作的均值和对数方差
        mean, log_std = self(state_v)
        # 得到标准差
        std = torch.exp(log_std)
        # 创建一个标准正态分布
        normal = dist.Normal(0, 1)
        # 从标准正态分布中采样和mean一样的维度的动作噪声
        z = normal.sample(mean.shape).to(device)
        # todo tanh是为了将动作限制在-1和1之间
        # mean + std * z在数学上的作用是将从标准正太分布转换为均值为mean，标准差为std的分布
        # 所以这里是将从标准正太分布中采样的动作噪声转换为预测的动作分布
        # 为什么不直接构建一个均值为mean，标准差为std的正态分布呢？
        # 您提出了一个非常好的问题。这种方法被称为"重参数化技巧"（Reparameterization Trick），是现代深度强化学习和变分推断中的一个关键技术。让我详细解释为什么我们使用这种方法：
        #
        #
        #
        # ```python
        # import torch
        # import torch.distributions as dist
        #
        # # 方法 1: 直接从指定均值和标准差的分布采样
        # def direct_sampling(mean, std):
        #     distribution = dist.Normal(mean, std)
        #     action = distribution.sample()
        #     return action
        #
        # # 方法 2: 使用重参数化技巧
        # def reparameterized_sampling(mean, std):
        #     normal = dist.Normal(0, 1)
        #     z = normal.sample(mean.shape)
        #     action = mean + std * z
        #     return action
        #
        # # 示例使用
        # mean = torch.tensor([0.5, -0.5])
        # std = torch.tensor([0.1, 0.2])
        #
        # action1 = direct_sampling(mean, std)
        # action2 = reparameterized_sampling(mean, std)
        #
        # print("Direct sampling:", action1)
        # print("Reparameterized sampling:", action2)
        #
        # # 梯度计算示例
        # mean.requires_grad_(True)
        # std.requires_grad_(True)
        #
        # action_reparam = reparameterized_sampling(mean, std)
        # loss = action_reparam.sum()
        # loss.backward()
        #
        # print("Gradient of mean:", mean.grad)
        # print("Gradient of std:", std.grad)
        #
        # ```
        #
        # 现在，让我解释为什么我们使用重参数化技巧，而不是直接采样：
        #
        # 1. 梯度传播：
        #    - 直接采样：如果我们直接从 N(mean, std) 采样，采样操作会阻断梯度的反向传播。这是因为采样操作本身不是可微的。
        #    - 重参数化：通过将随机性分离到独立的标准正态分布中，我们创建了一个可微的路径从输出（action）到输入参数（mean 和 std）。
        #
        # 2. 降低方差：
        #    - 重参数化技巧通常会产生较低方差的梯度估计，这有助于更稳定的训练过程。
        #
        # 3. 可控的随机性：
        #    - 通过分离随机性（z）和确定性部分（mean 和 std），我们可以更好地控制探索过程。
        #
        # 4. 数学等价性：
        #    - 从数学上讲，`mean + std * z`（其中 z ~ N(0, 1)）等价于从 N(mean, std^2) 直接采样。但前者允许梯度流动。
        #
        # 5. 便于优化：
        #    - 在训练过程中，我们可以直接优化 mean 和 std，而不是试图优化一个完整的分布。
        #
        # 6. 适用于更复杂的分布：
        #    - 这种技巧可以扩展到更复杂的分布，不仅限于高斯分布。
        #
        # 7. 一致性：
        #    - 在评估和推理时，我们可以通过设置 z = 0 来获得确定性的行为，这与训练时的随机行为保持一致。
        #
        # 8. 计算效率：
        #    - 在某些情况下，重参数化可能比直接从复杂分布采样更高效。
        #
        # 9. 便于实现高级技巧：
        #    - 例如，在SAC中，我们可以很容易地在采样后应用 tanh 函数来约束动作范围。
        #
        # 实际影响：
        # - 训练稳定性：使用重参数化技巧通常会导致更稳定的训练过程。
        # - 学习效率：通常可以更快地学习到有效的策略。
        # - 探索-利用平衡：提供了一种自然的方式来平衡探索和利用。
        #
        # 在SAC中的应用：
        # 在SAC算法中，重参数化技巧不仅用于生成动作，还用于计算策略的熵。这使得算法能够同时优化期望回报和策略熵，从而实现高效的探索和鲁棒的学习。
        #
        # 总结：
        # 重参数化技巧是一个强大的工具，它允许我们在保持随机性的同时，使得采样过程可微分。这在基于梯度的优化方法（如深度强化学习中的策略梯度方法）中特别有用。它解决了直接从参数化分布采样时梯度无法流动的问题，使得端到端的训练成为可能。
        action = torch.tanh(mean + std * z)
        # action = action * self.action_range
        # 动作的对数概率 以下计算包含三个部分：
        # dist.Normal(mean, std).log_prob(mean + std * z)
        # 计算原始高斯分布的对数概率，也可以理解为高斯分布下的动作概率
        # 构建一个均值为mean，标准差为std的高斯分布，然后计算mean + std * z的对数概率（动作概率）
        # 也就是对之前采样得到的动作z进行概率计算
        # torch.log(1. - action ** 2 + epsilon)
        # tanh 变换导致的概率密度变化的修正，由于预测的动作经过了tanh函数的变换，概率分布出现了变化，所以需要
        # 进行修正，这里的epsilon是为了数值稳定性，数学知识：雅可比行列式
        # torch.log(self.action_range_tensor)
        # 根据论文原文是对动作的范围进行修正，标准化不同的动作范围的动作概率
        log_prob = dist.Normal(mean, std).log_prob(mean + std * z) - torch.log(1. - action ** 2 + epsilon) - torch.log(self.action_range_tensor)
        # 这里是将预测的动作每一个维度的概率求和，得到一个综合的动作概率，因为在预测是，会对动作的每一个维度的概率进行预测
        log_prob = log_prob.sum(dim=1).unsqueeze(1)

        return action, log_prob, z, mean, log_std


    def get_action(self, state, device, greedy=False):
        '''
        该方法主要用于采样阶段
        :param state:
        :param device:
        :param greedy:
        :return:
        '''

        # 与TD3的不同点，对于获取的动作探索
        # 采用的均值和方差的计算方式
        state_v = state
        # 获取均值和对数标准差
        mean, log_std = self.forward(state_v)
        mean, log_std = mean.data.cpu(), log_std.data.cpu()
        # 将对数标准差转换为标准差
        std = np.exp(log_std)
        # 构建标准正太分布，并从中采样和预测的动作均值一样的维度的动作噪声
        normal = dist.Normal(0, 1)
        z = normal.sample(mean.shape)
        # 将噪声添加仅预测的动作均值中，生成动作
        # 并将值压缩到[-1, 1]，在转换为实际游戏要执行的动作范围
        # 从数学上来说，mean + std * z是将上面均值为0标准差为1的噪声分布转换为
        # 预测的动作均值mean和标准差std的分布中，也就是对动作进行了扰动，也就是探索
        # todo 这样做之后，虽然进行了随机探索，但是由于公式的是明确的，所以梯度是可以传播的
        # 整个表达式是可微的，可以用来训练网络
        action = torch.tanh(mean + std * z)
        action = action * self.action_range

        # 如果在测试、验证过程中，需要将greedy设置为True
        # 这个时候就不添加噪声，而是直接输出mean
        action = self.action_range * torch.tanh(mean) if greedy else action
        return action.cpu().numpy()


    def sample_action(self):
        action = torch.rand((1, self.act_size)) * 2 - 1
        return (self.action_range * action).cpu().numpy()


class TD3Actor(nn.Module):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, obs_size, act_size, action_range = 1):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(TD3Actor, self).__init__()

        self.net = nn.Sequential(
            nn.Linear(obs_size, 400),
            nn.ReLU(),
            nn.Linear(400, 300),
            nn.ReLU(),
            nn.Linear(300, act_size),
            nn.Tanh()
        )

        self.action_range = action_range

    def forward(self, x):
        return self.net(x) * self.action_range

class DDPGActorRGB(nn.Module):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGActorRGB, self).__init__()

        obs_size = (obs_size[2], obs_size[0], obs_size[1])
        self.net = nn.Sequential(
            nn.Conv2d(obs_size[0], 32, kernel_size=8, stride=4),
            nn.LayerNorm([32, 19, 19]),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.LayerNorm([64, 8, 8]),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.LayerNorm([64, 6, 6]),
            nn.ReLU()
        )

        conv_out_size = self._get_conv_out(obs_size)
        self.fc = nn.Sequential(
            nn.Linear(conv_out_size, 256),
            nn.ReLU(),
            nn.Linear(256, act_size),
            nn.Tanh()
        )

    def _get_conv_out(self, shape):
        o = self.net(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        fx = x.float() / 255.0

        return 2 * self.fc(self.net(fx).view(fx.size()[0], -1))


class DDPGCriticRGB(nn.Module):
    '''
    深度确定性策略梯度网络Q值评价网络
    '''
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGCriticRGB, self).__init__()

        obs_size = (obs_size[2], obs_size[0], obs_size[1])
        self.net = nn.Sequential(
            nn.Conv2d(obs_size[0], 32, kernel_size=8, stride=4),
            nn.LayerNorm([32, 19, 19]),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.LayerNorm([64, 8, 8]),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.LayerNorm([64, 6, 6]),
            nn.ReLU()
        )

        conv_out_size = self._get_conv_out(obs_size)

        # 构建观察+动作合并网路，输出评价Q值
        # 以前的Q值网络是输出每个动作的Q值，在连续值里面是直接输出评价Q值
        # 这里的400是self.obs_net的输出维度，act_size是动作网络的输出维度
        # 为了后续合并预测评价做准备
        self.out_net = nn.Sequential(
            nn.Linear(conv_out_size + act_size, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

    def _get_conv_out(self, shape):
        o = self.net(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x, a):
        fx = x.float() / 255.0
        conv_out = self.net(fx).view(fx.size()[0], -1)
        return self.out_net(torch.cat([conv_out, a], dim=1))
    

def bottle(f, x_tuple):
    x_sizes = tuple(map(lambda x: x.size(), x_tuple))
    y = f(*map(lambda x: x[0].contiguous().view(x[1][0] * x[1][1], *x[1][2:]), zip(x_tuple, x_sizes)))
    y_size = y.size()
    output = y.view(x_sizes[0][0], x_sizes[0][1], *y_size[1:])
    return output


class Dreamer(nn.Module):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, belief_size, state_size, act_size, hidden_size, embedding_size, min_std_dev=0.1):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(Dreamer, self).__init__()

        self.min_std_dev = min_std_dev

        self.fc_embed_state_action = nn.Linear(state_size + act_size, belief_size)
        self.rnn = nn.GRUCell(belief_size, belief_size)
        self.fc_embed_belief_prior = nn.Linear(belief_size, hidden_size)
        self.fc_state_prior = nn.Linear(hidden_size, 2 * state_size)
        self.fc_embed_belief_posterior = nn.Linear(belief_size + embedding_size, hidden_size)
        self.fc_state_priorior = nn.Linear(hidden_size, 2 * state_size)



    def forward(self, prev_state, actions, prev_belief, observations: Optional[torch.Tensor] = None, nonterminals: Optional[torch.Tensor] = None):
        T = actions.size(0) + 1
        beliefs, prior_states, prior_means, prior_std_devs, posterior_states, posterior_means, posterior_std_devs = (
            [torch.empty(0)] * T,
            [torch.empty(0)] * T,
            [torch.empty(0)] * T,
            [torch.empty(0)] * T,
            [torch.empty(0)] * T,
            [torch.empty(0)] * T,
            [torch.empty(0)] * T
        )

        beliefs[0], prior_states[0], posterior_states[0] = prev_belief, prev_state, prev_state

        for t in range(T - 1):
            _state = (
                prior_states[t] if observations is None else posterior_states[t]
            )
            _state = (
                _state if nonterminals is None else _state * nonterminals[t]
            )

            hidden = F.elu(self.fc_embed_state_action(torch.cat([_state, actions[t]], dim=1)))
            beliefs[t+1] = self.rnn(hidden, beliefs[t])
            hidden = F.elu(self.fc_embed_belief_prior(beliefs[t+1]))
            prior_means[t + 1], _prior_std_dev = torch.chunk(self.fc_state_prior(hidden), 2, dim=1)
            prior_std_devs[t + 1] = F.softplus(_prior_std_dev) + self.min_std_dev
            prior_states[t + 1] = prior_means[t + 1] + prior_std_devs[t + 1] * torch.randn_like(prior_means[t + 1])
            if observations is not None:
                t_ = t - 1
                hidden = F.elu(self.fc_embed_belief_posterior(torch.cat([beliefs[t + 1], observations[t_ + 1]], dim=1)))
                posterior_means[t + 1], _posterior_std_dev = torch.chunk(self.fc_state_priorior(hidden), 2, dim=1)
                posterior_std_devs[t + 1] = F.softplus(_posterior_std_dev) + self.min_std_dev
                posterior_states[t + 1] = posterior_means[t + 1] + posterior_std_devs[t + 1] * torch.randn_like(posterior_means[t + 1])

        hidden = [
            torch.stack(beliefs[1:], dim=0),
            torch.stack(prior_states[1:], dim=0),
            torch.stack(prior_means[1:], dim=0),
            torch.stack(prior_std_devs[1:], dim=0),
        ]

        if observations is not None:
            hidden += [
                torch.stack(posterior_states[1:], dim=0),
                torch.stack(posterior_means[1:], dim=0),
                torch.stack(posterior_std_devs[1:], dim=0)
            ]

        return hidden
    

class LinearDreamerObservation(nn.Module):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, obs_size, state_size, belief_size, embedding_size, min_std_dev=0.1):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(LinearDreamerObservation, self).__init__()

        self.net = nn.Sequential(
            nn.Linear(belief_size + state_size, embedding_size),
            nn.ELU(),
            nn.Linear(embedding_size, embedding_size),
            nn.ELU(),
            nn.Linear(embedding_size, obs_size)
        )
        


    def forward(self, belief, state):
        return self.net(torch.cat([belief, state], dim=1))
    

class ConvDreamerObservation(nn.Module):

    def __init__(self, belief_size, state_size, embedding_size):
        super().__init__()

        self.fc1 = nn.Linear(belief_size + state_size, embedding_size)
        self.embedding_size = embedding_size
        self.conv_decoder = nn.Sequential(
            nn.ConvTranspose2d(embedding_size, 128, kernel_size=5, stride=2),
            nn.ELU(),
            nn.ConvTranspose2d(128, 64, kernel_size=5, stride=2),
            nn.ELU(),
            nn.ConvTranspose2d(64, 32, kernel_size=6, stride=2),
            nn.ELU(),
            nn.ConvTranspose2d(32, 3, kernel_size=6, stride=2),
        )


    def forward(self, belief, state):
        x = torch.cat([belief, state], dim=1)
        x = self.fc1(x)
        x = x.view(-1, self.embedding_size, 1, 1)
        return self.conv_decoder(x)


class RewardDreamerModel(nn.Module):

    def __init__(self, belief_size, state_size, hidden_size):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(belief_size + state_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 1)
        )

    
    def forward(self, belief, state):
        return self.net(torch.cat([belief, state], dim=1)).squeeze(1)
    

class LinearDreamerEncoder(nn.Module):

    def __init__(self, obs_size, embedding_size):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(obs_size, embedding_size),
            nn.ReLU(),
            nn.Linear(embedding_size, embedding_size),
            nn.ReLU(),
            nn.Linear(embedding_size, embedding_size)
        )

    
    def forward(self, x):
        return self.net(x)
    


class ConvDreamerEncoder(nn.Module):

    def __init__(self, embedding_size):

        self.embedding_size = embedding_size
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=4, stride=2),
            nn.ReLU()
        )

        self.fc = nn.Identity() if embedding_size == 1024 else nn.Linear(1024, embedding_size)


    def forward(self, x):
        x = self.conv1(x)
        x = x.view(-1, 1024)
        return self.fc(x)
    


class DreamerActorModel(nn.Module):

    def __init__(self, belief_size, state_size, hidden_size, act_size, dist='tanh_normal', min_std=1e-4, init_std=5, mean_scale=5, device='cpu'):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(belief_size + state_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 2 * act_size)
        )

        self._dist = dist
        self._min_std = torch.tensor(min_std, dtype=torch.float).to(device=device).detach()
        self._init_std = torch.tensor(init_std, dtype=torch.float).to(device=device).detach()
        self._mean_scale = torch.tensor(mean_scale, dtype=torch.float).to(device=device).detach()

    
    def forward(self, belief, state):
        raw_init_std = torch.log(torch.exp(self._init_std) - 1)
        x = torch.cat([belief, state], dim=1)
        action = self.net(x).squeeze(dim=1)

        action_mean, action_std_dev = torch.chunk(action, 2, dim=1)
        action_mean = self._mean_scale * torch.tanh(action_mean / self._mean_scale)
        action_std = F.softplus(action_std_dev + raw_init_std) + self._min_std

        return action_mean, action_std
    

    def get_action(self, belief, state, det=False):
        action_mean, action_std = self.forward(belief, state)
        dist = torch.distributions.normal.Normal(action_mean, action_std)
        dist = torch.distributions.transformed_distribution.TransformedDistribution(dist, TanhBijector())
        dist = torch.distributions.Independent(dist, 1)
        dist = SampleDist(dist)

        if det:
            return dist.mode()
        else:
            return dist.rsample()
        

class SampleDist:

    def __init__(self, dist, samples=100):
        self._dist = dist
        self._samples = samples

    
    @property
    def name(self):
        return 'SampleDist'
    

    def __getattr__(self, name):
        return getattr(self._dist, name)
    

    def mean(self):
        sample = self._dist.rsample()
        return torch.mean(sample, dim=0)
    

    def mode(self):
        dist = self._dist.expand((self._samples, *self._dist.batch_shape))
        sample = dist.rsample()
        logprob = dist.log_prob(sample)
        batch_size = sample.size(1)
        feature_size = sample.size(2)
        indices = torch.argmax(logprob, dim=0).reshape(1, batch_size, 1).expand(1, batch_size, feature_size)
        return torch.gather(sample, 0, indices).squeeze(0)
    

    def entropy(self):
        dist = self._dist.expand((self._samples, *self._dist.batch_shape))
        sample = dist.rsample()
        logprob = dist.log_prob(sample)
        return -torch.mean(logprob, dim=0)
    

    def sample(self):
        return self._dist.sample()
        

class TanhBijector(torch.distributions.Transform):

    def __init__(self):
        super().__init__()

        self.bijective = True
        self.domain = torch.distributions.constraints.real
        self.codomain = torch.distributions.constraints.interval(-1.0, 1.0)

    
    @property
    def sign(self):
        return 1.0
    

    @staticmethod
    def atanh(x):
        return 0.5 * torch.log((1 + x) / (1 - x))
    

    def _call(self, x):
        return torch.tanh(x)
    

    def _inverse(self, y: torch.Tensor):
        y = torch.where((torch.abs(y) <= 1.0), torch.clamp(y, -0.99999997, 0.99999997), y)
        y = TanhBijector.atanh(y)
        return y
    

    def log_abs_det_jacobian(self, x, y):
        return 2.0 * (np.log(2) - x - F.softplus(-2.0 * x))
    


class DreamerValueModel(nn.Module):
    def __init__(self, belief_size, state_size, hidden_size):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(belief_size + state_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, 1)
        )


    def forward(self, belief, state):
        return self.net(torch.cat([belief, state], dim=1)).squeeze(1)


class DDPGActor(nn.Module):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGActor, self).__init__()

        self.net = nn.Sequential(
            nn.Linear(obs_size, 400),
            nn.ReLU(),
            nn.Linear(400, 300),
            nn.ReLU(),
            nn.Linear(300, act_size),
            nn.Tanh()
        )

    def forward(self, x):
        return self.net(x)


class DDPGCritic(nn.Module):
    '''
    深度确定性策略梯度网络Q值评价网络
    '''
    def __init__(self, obs_size, act_size):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGCritic, self).__init__()

        # 构建环境观察网络
        self.obs_net = nn.Sequential(
            nn.Linear(obs_size, 400),
            nn.ReLU(),
        )

        # 构建观察+动作合并网路，输出评价Q值
        # 以前的Q值网络是输出每个动作的Q值，在连续值里面是直接输出评价Q值
        # 这里的400是self.obs_net的输出维度，act_size是动作网络的输出维度
        # 为了后续合并预测评价做准备
        self.out_net = nn.Sequential(
            nn.Linear(400 + act_size, 300),
            nn.ReLU(),
            nn.Linear(300, 1)
        )

    def forward(self, x, a):
        obs = self.obs_net(x)
        return self.out_net(torch.cat([obs, a], dim=1))
    

class DDPGKanActor(nn.Module):
    '''
    深度确定性策略梯度动作预测网络
    '''
    def __init__(self, obs_size, act_size, use_eff_kan = False ):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGKanActor, self).__init__()

        self.net = nn.Sequential(
            EFF_KAN([obs_size, 64 ,256]),
            EFF_KAN([256, 64, 128]),
            EFF_KAN([128, 64, act_size]),
            nn.Tanh()
        ) if use_eff_kan else nn.Sequential(
            LIN_KAN([obs_size, 64 ,256]),
            LIN_KAN([256, 64, 128]),
            LIN_KAN([128, 64, act_size]),
            nn.Tanh()
        )

    def forward(self, x):
        return self.net(x)


class DDPGKanCritic(nn.Module):
    '''
    深度确定性策略梯度网络Q值评价网络
    '''
    def __init__(self, obs_size, act_size, use_eff_kan = False ):
        '''
        obs_size: 环境的维度
        act_size: 能够同时执行动作的个数（比如有多个手，不是每个手可以执行哪些动作）
        '''
        super(DDPGKanCritic, self).__init__()

        # 构建环境观察网络
        self.obs_net = nn.Sequential(
            EFF_KAN([obs_size, 64, 256]),
        ) if use_eff_kan else nn.Sequential(
            LIN_KAN([obs_size, 64, 256])
        )

        # 构建观察+动作合并网路，输出评价Q值
        # 以前的Q值网络是输出每个动作的Q值，在连续值里面是直接输出评价Q值
        # 这里的400是self.obs_net的输出维度，act_size是动作网络的输出维度
        # 为了后续合并预测评价做准备
        self.out_net = nn.Sequential(
            EFF_KAN([256 + act_size, 64, 128]),
            EFF_KAN([128, 32, 1])
        ) if use_eff_kan else nn.Sequential(
            LIN_KAN([256 + act_size, 64, 128]),
            LIN_KAN([128, 32, 1])
        )

    def forward(self, x, a):
        obs = self.obs_net(x)
        return self.out_net(torch.cat([obs, a], dim=1))
    

class AgentDreamer(ptan.agent.BaseAgent):
    """
    Agent implementing Orstein-Uhlenbeck exploration process
    实现一个具备探索能力的智能体
    """
    def __init__(self, net, device="cpu"):
        self.net = net
        self.device = device

    def initial_state(self):
        # 统一接口，但是这里不需要使用，返回None而不是使用pass
        return None

    def __call__(self, states, agent_states):
        '''
        states：当前的环境状态
        agent_states: 内部智能体的状态,之前的代理器里面这个基本无用,一开始的时候，agent_states是空的，但在这里因为要使用OU过程对执行的动作进行噪音干扰，所以需要使用了智能体的内部状态
        '''
        # 将环境转换为目标的数据类型
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)
        # 得到执行的动作输出
        mu_v = self.net(states_v)
        actions = mu_v.data.cpu().numpy()

        actions = np.clip(actions, -2, 2)
        return actions, None


class AgentDDPG(ptan.agent.BaseAgent):
    """
    Agent implementing Orstein-Uhlenbeck exploration process
    实现一个具备探索能力的智能体
    """
    def __init__(self, net, device="cpu", ou_enabled=True, ou_mu=0.0, ou_teta=0.15, ou_sigma=0.2, ou_epsilon=1.0):
        '''
        net: 动作执行网络
        device: 执行的设备
        ou_enabled: 是否开启OU过程探索
        ou_mu: 查看OU过程说明
        ou_teta: 查看OU过程说明
        ou_sigma: 查看OU过程说明
        ou_epsilon: 用于控制OU过程对最终动作的执行影响大小

        Ornstein-Uhlenbeck（OU）过程是一个用来生成噪声的数学工具，特别适合用于模拟物理过程中存在的摩擦或阻力的场景，因此它在强化学习中被用作探索机制，尤其是在连续动作空间的问题中。

        在使用OU过程时，参数`ou_teta`、`ou_mu`、`ou_sigma`通常是经验性选择的，他们分别控制：

        1. **`ou_teta`（Theta）** - 这个参数决定了系统回到均值（或趋向平稳状态）的速度。较高的`ou_teta`值会使系统更快地返回到均值，即产生的噪声会更快地趋于稳定状态。如果`ou_teta`设置得太高，系统可能会过于迅速地稳定下来，从而减少探索。太低，则可能导致系统的探索过于缓慢和延迟。

        2. **`ou_mu`（Mu）** - 这是长期平均或均值项，OU过程将围绕此均值进行波动。在强化学习的上下文中，`ou_mu`通常设置为0，这意味着没有任何行动的偏好，确保探索不会偏向于任何特定的行动。

        3. **`ou_sigma`（Sigma）** - 这个参数代表了波动的振幅或标准差，决定了噪声的强度或大小。较大的`ou_sigma`增加了探索的范围和随机性，但可能导致智能体行为变得不稳定；较小的值将减少探索性噪声，可能导致智能体过早收敛到次优策略。

        在确定这些参数时，通常需要通过实验调整以找到适合特定任务的值。一些普遍的做法和建议包括：

        - 开始时使用中等范围的值进行初步测试。
        - 观察智能体的表现，如果它行动过于随机，可以降低`ou_sigma`。
        - 如果智能体似乎探索得不够，或者过早收敛到一个行动上，可以增加`ou_sigma`。
        - 调整`ou_teta`以控制噪声的“记忆”效果，即噪声持续影响行动的程度。

        调整这些参数通常需要反复试验，因为理想的值很大程度上取决于具体问题、智能体的体系结构、奖励结构以及训练过程的其他方面。此外，有时会随着训练的进展而逐步减少噪声的影响，这是通过减少`ou_epsilon`来实现的，这种做法类似于退火。

        OU过程中的参数值ou_mu、ou_theta、ou_sigma和ou_dt通常根据具体问题来确定,没有一个通用的设置方法。下面是一些确定这些参数的常见思考:

        1. ou_mu - 长期均值可以设置为目标值或期望的平均奖励值。

        2. ou_theta - 回归速度可以先设置为一个较小的值(例如0.15),表示轻度回归到均值。可以通过调参来获得更好的回归速度。

        3. ou_sigma - 随机扰动标准差可以先设置较小(例如0.3),然后通过调参确定合适的探索强度。

        4. ou_dt - 时间步长可以设置得较小(例如0.01),然后确保dt足够小以维持迭代稳定性。

        5. ou_x0 - 通常初始化为0即可。

        另外,这些参数还需要互相平衡。例如theta越大,则需要更大的sigma来增加探索。

        一个通用的思路是:

        - 先用较小的回归速度θ,较大的扰动σ开始,更偏向探索。

        - 然后观察过程的变化,如果变化太随机可以适当增大θ或减小σ。

        - 调节参数直到得到既能利用已知信息又能适度探索的平衡效果。

        所以确定OU过程的参数需要根据具体问题,通过试错来找到一个合适的组合。
        '''
        self.net = net
        self.device = device
        self.ou_enabled = ou_enabled
        self.ou_mu = ou_mu
        self.ou_teta = ou_teta
        self.ou_sigma = ou_sigma
        self.ou_epsilon = ou_epsilon

    def initial_state(self):
        # 统一接口，但是这里不需要使用，返回None而不是使用pass
        return None

    def __call__(self, states, agent_states):
        '''
        states：当前的环境状态
        agent_states: 内部智能体的状态,之前的代理器里面这个基本无用,一开始的时候，agent_states是空的，但在这里因为要使用OU过程对执行的动作进行噪音干扰，所以需要使用了智能体的内部状态
        '''
        # 将环境转换为目标的数据类型
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)
        # 得到执行的动作输出
        mu_v = self.net(states_v)
        actions = mu_v.data.cpu().numpy()

        if self.ou_enabled and self.ou_epsilon > 0:
            # 启动随机探索
            new_a_states = []
            for a_state, action in zip(agent_states, actions):
                if a_state is None:
                    # 为什么智能体状态为空时可以构建一个维度和动作相同，且全为零的状态
                    # 这是针对初始状态，初始情况下，智能体的状态为空，则创建一个为0的状态
                    a_state = np.zeros(shape=action.shape, dtype=np.float32)
                # 书p298也 todo 了解OU过程
                a_state += self.ou_teta * (self.ou_mu - a_state)
                a_state += self.ou_sigma * np.random.normal(size=action.shape)

                action += self.ou_epsilon * a_state
                new_a_states.append(a_state)
        else:
            new_a_states = agent_states

        # 这个步骤 的意思是修正动作值到-1和1之间，否则将导致PyBullet抛出异常
        actions = np.clip(actions, -1, 1)
        return actions, new_a_states
    

class AgentDDPGPendulum(ptan.agent.BaseAgent):
    """
    Agent implementing Orstein-Uhlenbeck exploration process
    实现一个具备探索能力的智能体
    """
    def __init__(self, net, device="cpu", ou_enabled=True, ou_mu=0.0, ou_teta=0.15, ou_sigma=0.2, ou_epsilon=1.0):
        self.net = net
        self.device = device
        self.ou_enabled = ou_enabled
        self.ou_mu = ou_mu
        self.ou_teta = ou_teta
        self.ou_sigma = ou_sigma
        self.ou_epsilon = ou_epsilon
        self.agent_states = None

    def initial_state(self):
        # 统一接口，但是这里不需要使用，返回None而不是使用pass
        return None

    def __call__(self, states, agent_states):
        '''
        states：当前的环境状态
        agent_states: 内部智能体的状态,之前的代理器里面这个基本无用,一开始的时候，agent_states是空的，但在这里因为要使用OU过程对执行的动作进行噪音干扰，所以需要使用了智能体的内部状态
        '''
        # 将环境转换为目标的数据类型
        states_v = ptan.agent.float32_preprocessor(states).to(self.device)
        # 得到执行的动作输出
        mu_v = self.net(states_v)
        actions = mu_v.data.cpu().numpy()

        if self.ou_enabled and self.ou_epsilon > 0:
            # 启动随机探索
            new_a_states = []
            for a_state, action in zip(agent_states, actions):
                if a_state is None:
                    # 为什么智能体状态为空时可以构建一个维度和动作相同，且全为零的状态
                    # 这是针对初始状态，初始情况下，智能体的状态为空，则创建一个为0的状态
                    a_state = np.zeros(shape=action.shape, dtype=np.float32)
                # 书p298也 todo 了解OU过程
                a_state += self.ou_teta * (self.ou_mu - a_state)
                a_state += self.ou_sigma * np.random.normal(size=action.shape)

                action += self.ou_epsilon * a_state
                new_a_states.append(a_state)
        else:
            new_a_states = agent_states

        actions = np.clip(actions, -2, 2)
        self.agent_states = new_a_states
        return actions, new_a_states
    
    def state_dict(self):
        return {
            'ou_enabled': self.ou_enabled,
            'ou_mu': self.ou_mu,
            'ou_teta': self.ou_teta,
            'ou_sigma': self.ou_sigma,
            'ou_epsilon': self.ou_epsilon,
            'agent_states': self.agent_states
        }

    def load_state_dict(self, state_dict):
        self.ou_enabled = state_dict['ou_enabled']
        self.ou_mu = state_dict['ou_mu']
        self.ou_teta = state_dict['ou_teta']
        self.ou_sigma = state_dict['ou_sigma']
        self.ou_epsilon = state_dict['ou_epsilon']
        self.agent_states = state_dict['agent_states']


class AgentPlaNet(ptan.agent.BaseAgent):
    """
    Agent implementing Orstein-Uhlenbeck exploration process
    实现一个具备探索能力的智能体
    """
    def __init__(self, planet, action_dim, params, device="cpu"):
        self.planet = planet
        self.device = device
        self.agent_states = None
        self.params = params
        self.action_dim = action_dim
        self.d_type = common.get_dtype(self.params['fp_precision'])
        self.action_epsilon = float(self.params['action_epsilon'])

    def initial_state(self):
        h_state = self.planet.get_init_h_state(batch_size=1)
        return h_state
    

    def __call__(self, states, agent_states):
        '''
        收集新的episode
        '''
        if agent_states is None or agent_states[0] is None:
            h_state = self.initial_state()[0]
        else:
            h_state = agent_states[0]
        # Inject observation noise
        # 这里的噪声是一个高斯噪声，均值为0，方差为1
        # 将噪声进行缩放后，使其范围满足和观察空间一致，因为高斯噪声本身就分布在0附近的左右，所以只需要使用正数缩放即可
        # 这行代码通过生成与 input_obs 同形状的标准正态分布噪声，并用 (1/2^(pixel_bit)) 缩放后加到 input_obs 上，从而为输入观测值注入噪声。这样做可以模拟低位深像素的量化效应或者作为数据正则化的一种手段。
        # 提高代码鲁棒性
        states = torch.from_numpy(np.array(states)).to(self.d_type).to(self.device)
        noisy_states = (1/pow(2, self.params['pixel_bit']))*torch.randn_like(states) + states
        # Get posterior states using observation 获取编码后的观察
        encoded_obs = self.planet.obs_encoder(noisy_states.to(self.device))
        # 根据上一个隐藏状态和编码后的观察生成一个后验分布
        posterior_z = self.planet.repr_model(h_state, encoded_obs)
        # 对后验分布进行采样，得到一个潜在状态
        # shape = （batch_size， z_dim）
        z_state = posterior_z.sample()
        # Get best action by planning in latent space through open-loop prediction
        with torch.no_grad():
            # 预测动作
            action = self.plan_action_with_cem(h_state, z_state)
        exploration_noise = Normal(loc=0.0, scale=self.action_epsilon).sample(sample_shape=torch.Size(action.shape)).to(self.d_type).to(self.device)
        noisy_action = action + exploration_noise # 对动作加入噪声，这里可以看出其针对连续动作空间，对离散动作空间不支持
        # Get next latent state 预测下一个确定性隐藏状态
        h_state = self.planet.rnn_model(h_state, z_state, noisy_action.unsqueeze(dim=0))

        return [action.cpu().numpy()], [h_state]

    
    def plan_action_with_cem(self, init_h_state, init_z_state):
        '''
        param init_h_state: 上一个隐藏状态
        param init_z_state: 本次潜在状态，也就是后验分布的潜在状态

        return: 根据不断的迭代得到认为的最好的动作
        '''
        # todo `planning_horizon` 参数指定了在规划过程中，评估候选动作序列时所预见的未来时间步数。在基于 Planet 算法的实现中，这个值控制了模型在每次规划时向前预测多少步，以便选择出最优的动作计划。
        action_dist = Independent(Normal(loc=torch.zeros(self.params['planning_horizon'], self.action_dim), scale=1.0), reinterpreted_batch_ndims=2)
        # todo plan_optimization_iter参数指定了在规划过程中，通过多次迭代优化候选动作序列，以便搜索出最优计划的迭代次数。也就是说，规划器会在10次迭代中不断改进候选计划
        for _ in range(self.params['plan_optimization_iter']):
            reward_buffer = list()
            # n_plans 数指定了在每次优化迭代中将采样多少条候选动作序列（计划）。也就是说，在每次规划过程中，会随机采样 1000 个动作序列，然后通过后续的评估机制选择出表现最佳的候选计划，以便在环境中执行
            # 也就是说随机生成n_plans个状态分布进行选择评估？
            h_state = torch.clone(init_h_state).repeat(self.params['n_plans'], 1)
            z_state = torch.clone(init_z_state).repeat(self.params['n_plans'], 1)
            # 生成候选动作，shape = todo
            candidate_plan = torch.clip_(
                action_dist.sample(sample_shape=torch.Size([self.params['n_plans']])).to(self.d_type).to(self.device),
                min=self.params['min_action'], max=self.params['max_action'])
            # 对每一个时间步的动作进行评估，这个时间步包含对未来观察的预测吧
            for time_step in range(self.params['planning_horizon']):
                batched_ts_action = candidate_plan[:, time_step, :]
                # Use learnt dynamics to get next hidden state 根据动作得到下一次的确定性隐藏状态
                h_state = self.planet.rnn_model(h_state, z_state, batched_ts_action)
                # Get latent variables from transition model (prior) 得到下一个时间先验状态
                prior_z = self.planet.transition_model(h_state)
                z_state = prior_z.sample() # 先验状态模拟后验状态
                # 利用确定性状态+模拟本地后验状态预测奖励分布
                predicted_reward = self.planet.reward_model(h_state, z_state)
                # 对奖励进行裁剪，因为可能存在跳帧，则认为跳帧的每次动作的到的奖励都一样，那么跳帧期间的奖励=单帧的奖励*次数
                sampled_reward = torch.clip(predicted_reward.mean,
                                            min=self.params['min_reward'], max=(1+self.params['action_repeat'])*self.params['max_reward'])
                # 保存奖励缓冲区
                reward_buffer.append(sampled_reward)
            # 展平所有时间步的奖励
            plan_reward = torch.stack(reward_buffer).squeeze().sum(dim=0)
            # 根据奖励的大小，选择topK个奖励最高的动作分布，得到新的动作分布和均值
            # 再创建新的动作分布，重新进行动作采样重新进行评估
            chosen_actions = candidate_plan[torch.topk(plan_reward, k=self.params['top_k']).indices]
            action_mu, action_std = chosen_actions.mean(dim=0), chosen_actions.std(dim=0)
            action_dist = Independent(Normal(loc=action_mu, scale=action_std+1e-6), reinterpreted_batch_ndims=2)
        optimized_next_action = action_dist.mean[0]
        return optimized_next_action
    

class RecurrentModel(nn.Module):
    def __init__(self, params, action_dim):
        '''
        params: 整个训练的参数
        action_dim: 动作的维度

        网络结构比较简单，只是一个 GRU 网络，输入维度为 z_dim + action_dim，输出维度为 h_dim。
        '''
        super(RecurrentModel, self).__init__()
        self.params = params
        # GRU 网络的输入维度为 z_dim + action_dim
        self.gru_net = nn.GRUCell(input_size=self.params['z_dim']+action_dim, hidden_size=self.params['h_dim'])

    def forward(self, h_state, z_state, action):
        gru_input = torch.concat([z_state, action], dim=1)
        next_h_state = self.gru_net(gru_input, h_state)
        return next_h_state
    


class EncoderModel(nn.Module):
    def __init__(self, obs_shape, params):
        '''
        parameters: 整个训练的参数

        网络结构比较简单，只是一个 CNN 网络，输入维度为 3（RGB 图像），输出维度为 256。
        输入3 * 64 * 64 ，输出 256 * 2 * 2。
        '''
        super(EncoderModel, self).__init__()
        self.params = params
        self.encoder_net = nn.Sequential(
            nn.Conv2d(in_channels=obs_shape[0], out_channels=32, kernel_size=4, stride=2, padding='valid'),
            nn.ReLU(),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2, padding='valid'),
            nn.ReLU(),
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding='valid'),
            nn.ReLU(),
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding='valid'),
            nn.ReLU(),
        )

    def forward(self, obs):
        encoded_obs = self.encoder_net(obs)
        # 这行代码使用了 [einops](https://einops.rocks/) 库中的 `rearrange` 函数，将原本形状为 `[batch, channel, height, width]` 的张量展开（flatten）成形状为 `[batch, channel * height * width]` 的张量。这通常用于将卷积层的输出扁平化，以便后续通过全连接层（或其他处理）使用。
        encoded_obs = eop.rearrange(encoded_obs, 'b c h w -> b (c h w)')
        return encoded_obs
    

class RepresentationModel(nn.Module):
    def __init__(self, params):
        '''
        parameters: 整个训练的参数

        网络结构比较简单，只是一个全连接网络，输入维度为 h_dim + feat_dim，输出维度为 2 * z_dim。
        输出的是一个正态分布的均值和标准差，表示当前时刻的潜在状态。
        '''
        super(RepresentationModel, self).__init__()
        self.params = params
        self.repr_net = FeedForwardNet(
            input_dim=self.params['h_dim']+self.params['feat_dim'],
            output_dim=2*self.params['z_dim'],
            hidden_dim=self.params['h_dim'],
            n_layers=self.params['n_ff_layers']
        )
        self.min_std = float(self.params['min_std'])

    def forward(self, h_state, encoded_obs):
        concat_input = torch.concat([h_state, encoded_obs], dim=1)
        mu, pre_std = torch.chunk(self.repr_net(concat_input), chunks=2, dim=1)
        # 对 pre_std 进行 softplus 操作并加上一个最小标准差，以确保标准差为正。
        std = F.softplus(pre_std + 0.55) + self.min_std
        dist_posterior = Independent(Normal(loc=mu, scale=std), reinterpreted_batch_ndims=1)
        return dist_posterior
    

class TransitionModel(nn.Module):
    def __init__(self, params):
        '''
        parameters: 整个训练的参数

        网络结构比较简单，只是一个全连接网络，输入维度为 h_dim，输出维度为 2 * z_dim。
        输出的是一个正态分布的均值和标准差，表示当前时刻的潜在状态。
        '''
        super(TransitionModel, self).__init__()
        self.params = params
        self.transition_net = FeedForwardNet(
            input_dim=self.params['h_dim'],
            output_dim=2*self.params['z_dim'],
            hidden_dim=self.params['h_dim'],
            n_layers=self.params['n_ff_layers']
        )
        self.min_std = float(self.params['min_std'])

    def forward(self, h_state):
        mu, pre_std = torch.chunk(self.transition_net(h_state), chunks=2, dim=1)
        # 对 pre_std 进行 softplus 操作并加上一个最小标准差，以确保标准差为正。
        std = F.softplus(pre_std + 0.55) + self.min_std
        dist_prior = Independent(Normal(loc=mu, scale=std), reinterpreted_batch_ndims=1)
        return dist_prior
    

class DecoderModel(nn.Module):
    def __init__(self, obs_shape, params):
        '''
        params : 整个训练的参数

        网络有一个全连接层，将 h_state 和 z_state 连接起来，然后通过一系列的反卷积层将其解码为与环境观测相对应的重构输出。
        也就是输出为 3 * 64 * 64 的大小的分布
        '''
        super(DecoderModel, self).__init__()
        self.params = params
        self.fc_net = nn.Linear(in_features=self.params['h_dim']+self.params['z_dim'], out_features=1024)
        self.decoder_net = nn.Sequential(
            nn.ConvTranspose2d(in_channels=256, out_channels=128, stride=2, kernel_size=5, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(in_channels=128, out_channels=64, stride=2, kernel_size=5),
            nn.ReLU(),
            nn.ConvTranspose2d(in_channels=64, out_channels=32, stride=2, kernel_size=6),
            nn.ReLU(),
            nn.ConvTranspose2d(in_channels=32, out_channels=obs_shape[0], stride=2, kernel_size=6)
        )

    def forward(self, h_state, z_state):
        concat_input = torch.concat([h_state, z_state], dim=-1)
        fc_output = self.fc_net(concat_input)
        reshaped_input = fc_output.view(-1, 256, 2, 2)
        mu_obs = self.decoder_net(reshaped_input)
        dist_obs = Independent(Normal(loc=mu_obs, scale=1.0), reinterpreted_batch_ndims=3)
        return dist_obs
    

class RewardModel(nn.Module):
    def __init__(self, params):
        '''
        params: 整个训练的参数

        网络结构比较简单，只是一个全连接网络，输入维度为 h_dim + z_dim，输出维度为 1。应该是直接预测奖励值。
        '''
        super(RewardModel, self).__init__()
        self.params = params
        self.reward_net = FeedForwardNet(
            input_dim=params['h_dim']+params['z_dim'],
            output_dim=1,
            hidden_dim=params['h_dim'],
            n_layers=self.params['n_ff_layers']
        )

    def forward(self, h_state, z_state):
        concat_input = torch.concat([h_state, z_state], dim=-1)
        mu_reward = self.reward_net(concat_input)
        dist_reward = Independent(Normal(loc=mu_reward, scale=1.0), reinterpreted_batch_ndims=1)
        return dist_reward
    
    
class FeedForwardNet(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim, n_layers):
        '''
        前向传播网络

        param input_dim: 输入维度
        param output_dim: 输出维度
        param hidden_dim: 隐藏层维度
        param n_layers: 隐藏层的层数
        
        '''

        super(FeedForwardNet, self).__init__()
        self.to_hidden = nn.Sequential(
            nn.Linear(in_features=input_dim, out_features=hidden_dim),
            nn.ReLU()
        )
        self.hidden = nn.Sequential(*[
            nn.Sequential(
                nn.Linear(in_features=hidden_dim, out_features=hidden_dim),
                nn.ReLU()
            ) for _ in range(n_layers-1)
        ])
        self.from_hidden = nn.Sequential(
            nn.Linear(in_features=hidden_dim, out_features=output_dim)
        )

    def forward(self, x):
        to_hidden = self.to_hidden(x)
        hidden = self.hidden(to_hidden)
        return self.from_hidden(hidden)


class Planet(nn.Module):
    def __init__(self, params, obs_shape, action_dim, device='cpu'):
        '''
        parameters: 整个训练的参数
        action_dim: 动作的维度
        '''

        super(Planet, self).__init__()
        self.params = params
        self.d_type = common.get_dtype(self.params['fp_precision'])
        self.device = device
        self.action_dim = action_dim
        '''
        这些模型在 PlaNet 框架中各有分工，简要说明如下：

        1. **EncoderModel**: 负责将观察（像素图像）编码为紧凑的特征向量，减少原始输入的维度。  
        2. **RepresentationModel**: 使用编码后的特征向量与先前的隐藏状态来生成后验分布 (posterior)，表示当前时刻的潜在状态。  
        3. **TransitionModel**: 生成先验分布 (prior)，描述在未观测到当前帧情况下，对潜在状态的预测。  
        4. **DecoderModel**: 将隐藏状态与潜在状态解码回与环境观测相对应的重构输出。  
        5. **RewardModel**: 从隐藏状态和潜在状态预测当前时刻的奖励，以便在学习及规划过程中使用。  
        6. **RecurrentModel**: 按序列更新并维护隐藏状态，结合潜在状态和动作来捕捉时间上的依赖关系。

        todo 描述这几个网络之间的数据流通情况
        '''
        self.rnn_model = RecurrentModel(params=self.params, action_dim=self.action_dim)
        self.obs_encoder = EncoderModel(obs_shape=obs_shape, params=self.params)
        self.repr_model = RepresentationModel(params=self.params)
        self.transition_model = TransitionModel(params=self.params)
        self.decoder_model = DecoderModel(params=self.params, obs_shape=obs_shape)
        self.reward_model = RewardModel(params=self.params)

    def __repr__(self):
        return 'PlaNet'

    def get_init_h_state(self, batch_size):
        '''
        创建一个初始的确定性隐藏状态
        '''
        return torch.zeros((batch_size, self.params['h_dim']), dtype=self.d_type, device=self.device)

    def forward(self, sampled_episodes):
        '''
        param sampled_episodes: 采样的 episode 数据
        sampled_episodes['obs']: 采样的观察数据chunk_length， batch_size， visual_resolution， visual_resolution， channels）
        sampled_episodes['action']: 采样的动作数据（chunk_length， batch_size， action_dim）
        sampled_episodes['reward']: 采样的奖励数据  chunk_length， batch_size， 1）
        '''
        # 创建一个字典来存储预测的分布
        # prior: 先验分布
        # posterior: 后验分布
        # recon_obs: 重构观察
        # reward: 奖励分布
        # 这里的分布是一个分布对象，包含均值和方差等信息 todo 
        dist_predicted = {'prior': list(), 'posterior': list(), 'recon_obs': list(), 'reward': list()}
        # 获取一个初始的确定性隐藏状态，shape 为（batch_size， h_dim）
        h_state = self.get_init_h_state(batch_size=self.params['batch_size'])
        for time_stamp in range(self.params['chunk_length']):
            # 遍历每一个序列的时间步
            input_obs = sampled_episodes['obs'][time_stamp] # 获取当前时间步的观察数据
            # 创建一个噪声观察数据，噪声的大小与像素位数有关
            # 这里的噪声是一个高斯噪声，均值为0，方差为1
            # 将噪声进行缩放后，使其范围满足和观察空间一致，因为高斯噪声本身就分布在0附近的左右，所以只需要使用正数缩放即可
            # 这行代码通过生成与 input_obs 同形状的标准正态分布噪声，并用 (1/2^(pixel_bit)) 缩放后加到 input_obs 上，从而为输入观测值注入噪声。这样做可以模拟低位深像素的量化效应或者作为数据正则化的一种手段。
            # 提高代码鲁棒性
            noisy_input_obs = (1/pow(2, self.params['pixel_bit']))*torch.randn_like(input_obs) + input_obs
            # 获取当前时间步的动作数据
            action = sampled_episodes['action'][time_stamp]

            # 对当前时间步的观察数据进行编码，得到一个编码后的观察数据
            # shape = （batch_size， h_dim）
            encoded_obs = self.obs_encoder(noisy_input_obs)
            # 根据上一个隐藏状态生（上个确定性隐藏状态（上个后验状态、上一个动作））成一个先验分布，shape = （batch_size， 2 * z_dim）
            z_prior = self.transition_model(h_state)
            # 根据编码后的观察数据和上一个隐藏状态生成一个后验分布（上个确定性隐藏状态（上个后验状态、上一个动作）， 本次编码后的观察），shape = （batch_size， 2 * z_dim）
            z_posterior = self.repr_model(h_state, encoded_obs)
            
            # 对后验分布进行采样，得到一个潜在状态
            # shape = （batch_size， z_dim）
            z_state = z_posterior.rsample()

            # 对潜在状态进行解码，得到一个重构观察数据，输入为潜在状态和上一个隐藏状态
            # shape = （batch_size， visual_resolution， visual_resolution， channels）
            dist_recon_obs = self.decoder_model(h_state, z_state)
            # 根据潜在状态和上一个隐藏状态生成一个奖励分布，输入为潜在状态和上一个隐藏状态
            # shape = （batch_size， 1）
            dist_reward = self.reward_model(h_state, z_state)
            # 根据潜在状态和动作数据更确定性新隐藏状态，输入为潜在状态和动作数据
            h_state = self.rnn_model(h_state, z_state, action)

            dist_predicted['prior'].append(z_prior)
            dist_predicted['posterior'].append(z_posterior)
            dist_predicted['recon_obs'].append(dist_recon_obs)
            dist_predicted['reward'].append(dist_reward)
        # 返回一个字典，包含了每个时间步先验分布，后验分布，重构观察数据和奖励分布
        return dist_predicted

    def compute_loss(self, target, dist_predicted):
        '''
        param target: 目标数据,也是采样的数据
        param dist_predicted: 预测的分布

        return net_loss: 总损失, 其他损失返回标量，用于记录
        '''
        # 创建一个大小和观察一致的分布
        sampled_reconstructed_obs = torch.stack([dist_recon_obs.rsample() for dist_recon_obs in dist_predicted['recon_obs']])
        # 创建一个大小和奖励一致的分布
        sampled_reward = torch.stack([dist_reward.rsample() for dist_reward in dist_predicted['reward']])
        # Individual loss terms
        # 重构观察损失，也就是重构观察数据和目标观察数据之间的均方误差
        recon_loss = ((target['obs'] - sampled_reconstructed_obs) ** 2).mean(dim=0).mean(dim=0).sum()
        # 预测中的先验分布和后验分布之间的 KL 散度损失，要保证其相近
        kl_loss = torch.stack(
            [dist.kl_divergence(p=dist_posterior, q=dist_prior) for dist_prior, dist_posterior in
             zip(dist_predicted['prior'], dist_predicted['posterior'])]
        )
        # 对 KL 散度损失进行平均
        # 定义了 KL 散度的下限，通常被称为“free nats”技巧。在计算 KL 散度损失时，将每个时间步的 KL 值与 free_nats 比较，确保即使模型预测的 KL 散度非常低，也不会过分惩罚。这样可以防止模型过度拟合噪声，提高训练的稳定性
        kl_loss = (torch.maximum(kl_loss, torch.tensor([self.params['free_nats']])[0])).mean()
        # 奖励预测损失，计算目标奖励数据和预测奖励数据之间的均方误差
        reward_prediction_loss = ((target['reward'] - sampled_reward) ** 2).mean()
        # Net loss term 计算总损失
        net_loss = recon_loss + kl_loss + reward_prediction_loss
        return net_loss, (recon_loss.item(), kl_loss.item(), reward_prediction_loss.item())